# __author__ = 'heyin'
# __date__ = '2018/11/13 9:00'
# 决策树的使用
import pandas as pd
from sklearn.tree import DecisionTreeClassifier, export_graphviz
from sklearn.feature_extraction import DictVectorizer
from sklearn.model_selection import train_test_split


def d_tree():
    """预测泰坦尼克生存率"""
    tt = pd.read_csv('titanic.csv')
    # 判断哪些列存在缺失值
    # ret = pd.isnull(tt)  # 得到的是bool类型矩阵，表示某个位置是否是nan
    # ret = ret.any()  # 得到每个列是否存在nan的series
    # print(ret)
    # print(ret.values)  # values用于取出df中的数据，取出二维数组
    # print(tt.keys())  # 取出列名，即 index，也是index类型，类似列表
    # tt = tt[pd.isnull(tt).values == True]  # 取出values中为真的数据，也就是包含了nan的数据
    # tt = tt.drop_duplicates()  # 一行元素全部相同时才去除，可以指定参数进行去重
    # print(tt)

    # 取出特征值和目标值，特征值取出pclass,age,sex  目标值 survived
    # age缺失值以平均值进行填充
    tt['age'][pd.isnull(tt['age'])] = tt['age'].mean(axis=0)
    x = tt[['pclass', 'sex', 'age']]
    y = tt['survived']

    # 特征抽取，应当是以字典的形式进行转换
    x = x.to_dict(orient='records')  # records,整体构成一个列表，内层是将原始数据的每行提取出来形成字典

    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25, random_state=1)
    print(len(x_train))
    dv = DictVectorizer(sparse=False)
    x_train = dv.fit_transform(x_train)
    x_test = dv.transform(x_test)

    dt = DecisionTreeClassifier()
    dt.fit(x_train, y_train)
    print(dv.get_feature_names())
    print(dt.feature_importances_)


    y_pred = dt.predict(x_test)
    print('测试集score', dt.score(x_test, y_test))
    print('训练集score', dt.score(x_train, y_train))
    export_graphviz(dt, out_file='./decisiontree/titanic.dot', feature_names=dv.get_feature_names())

if __name__ == '__main__':
    d_tree()
